# This file will be called by "MasterFile.R"
# remove(list=ls())
# source("/Users/wenhao/Dropbox/My Packages/package_function_and_path.R")
# setwd(dirname(rstudioapi::getActiveDocumentContext()$path))

# ------------------ Plot Specifications ---------------------
margins = c(3.5, 3.5, 1, 1)
plot_slippery_slope_two_economy = TRUE
plot_width = 4.2
plot_height = 3.8

# ----------------- Load the baseline calibrated parameters -------------
load(file = "baseline_calibration.Rdata")  # load the baseline calibrated parameters.
source("Model scripts.R")  # This script solves the model and provide simulation functions 
# IMPORTNAT: change eta. 
par$eta_H = 1.5*par$eta_H
par$eta_L = 1.5*par$eta_L
calibration_results=calibration_moments(par)
equi_sol = calibration_results$equi_sol
moments = cbind(calibration_results$key_moments, calibration_results$other_moments)

# Then assign values for next processing
g_bar = par$g_bar; d_bar = par$d_bar; dt = equi_sol$dt;
AH = par$AH;  AL = par$AL
lambda=par$lambda; beta=par$beta; r = par$r; eta_H=par$eta_H; eta_L=par$eta_L
w_grid = equi_sol$w_grid
f_list = calibration_results$f_list
solve_individual_optimal_faster = f_list$solve_individual_optimal_faster; F_fun = f_list$F_fun; Delta_output_fun = f_list$Delta_output_fun; solve_equilibrium=f_list$solve_equilibrium; Simulate_model=f_list$Simulate_model; solve_other_moments=f_list$solve_other_moments; 
muK_fun=f_list$muK_fun; DeltaK_fun=f_list$DeltaK_fun; muw_fun=f_list$muw_fun; Deltaw_fun=f_list$Deltaw_fun; investment_cost_fun=f_list$investment_cost_fun;  l_dN_fun=f_list$l_dN_fun;  post_crisis_fraction_of_capital=f_list$post_crisis_fraction_of_capital; output_drop_and_gbar=f_list$output_drop_and_gbar;  
solve_welfare = f_list$solve_welfare
simulation_fun=f_list$simulation_fun
avg_zeta = par$avg_zeta; 
qL = equi_sol$qL; qH = equi_sol$qH;  


#---------------- Figure A7: Government Financing and Crisis Dynamics -------------
g_bar_vec = seq(0.01, AH , length.out=50)
init_vec = rep(0, length(g_bar_vec))
TB = data.table( g_bar=g_bar_vec, qH=init_vec, qL=init_vec, iH=init_vec, iL=init_vec, creditor_profit_H_frac=init_vec, creditor_profit_L_frac=init_vec, 
                 H_zeta_ = init_vec, L_zeta_ = init_vec, dK=init_vec, dw=init_vec,
                 muw = init_vec, muw_middle=init_vec, muK=init_vec, muK_middle=init_vec, w_avg=init_vec,
                 avg_creditor_recovery = init_vec, avg_investment_to_GDP=init_vec, avg_output_drop=init_vec, takeup_ratio=init_vec,
                 strategic_bankruptcy=init_vec, liquidity_bankruptcy=init_vec, total_bankruptcy=init_vec, 
                 H_zeta_bar = init_vec, L_zeta_bar = init_vec,
                 r1 = init_vec,  r2=init_vec, d_hat1 = init_vec, d_hat2 = init_vec, 
                 xL1 = init_vec, xL2 = init_vec, xH1=init_vec, xH2=init_vec, output_change=init_vec
                 )
zeta_cases = c(0.1, 0.2)   # These two cases map to the interest rate r1 and r2. 
w0=calibration_results$equi_sol$w_steady   # Evaluate at the average w0
for(iter in 1:length(g_bar_vec)){
  print( paste0("Progress:" , round(iter/length(g_bar_vec)*100) , "%") )
  equi_sol = list()
  equi_sol = solve_equilibrium(d_bar, g_bar_vec[iter])
  equi_sol = solve_other_moments(equi_sol)
  w0 = equi_sol$w_avg
  TB$qH[iter] = equi_sol$qH
  TB$qL[iter] = equi_sol$qL
  TB$iH[iter] = equi_sol$iH
  TB$iL[iter] = equi_sol$iL
  TB$creditor_profit_H_frac[iter] = equi_sol$creditor_profit_frac_H
  TB$creditor_profit_L_frac[iter] = equi_sol$creditor_profit_frac_L
  TB$H_zeta_[iter] = equi_sol$zeta__H
  TB$L_zeta_[iter] = equi_sol$zeta__L
  TB$dK[iter] = DeltaK_fun( equi_sol$kappaH, equi_sol$kappaL, w0 )
  TB$dw[iter] = Deltaw_fun( equi_sol$kappaH, equi_sol$kappaL, w0 )
  TB$w_avg[iter] = equi_sol$w_avg
  TB$muw[iter] = muw_fun( equi_sol$iH, equi_sol$iL, w0 )
  TB$muw_middle[iter] = muw_fun( equi_sol$iH, equi_sol$iL, 0.5 )
  TB$muK[iter] = muK_fun( equi_sol$iH, equi_sol$iL, w0 )
  TB$muK_middle[iter] = muK_fun( equi_sol$iH, equi_sol$iL, 0.5 )
  TB$avg_creditor_recovery[iter] = equi_sol$avg_creditor_recovery
  TB$avg_investment_to_GDP[iter] = equi_sol$avg_investment_to_GDP
  TB$avg_output_drop[iter] = equi_sol$avg_output_drop
  TB$avg_productivity[iter] = equi_sol$avg_productivity
  TB$takeup_ratio[iter] = equi_sol$takeup_ratio
  TB$strategic_bankruptcy[iter] = equi_sol$strategic_bankruptcy_total
  TB$liquidity_bankruptcy[iter] = equi_sol$liquidity_bankruptcy_total
  TB$total_bankruptcy[iter]     = equi_sol$bankruptcy_total
  TB$H_zeta_bar[iter] = equi_sol$zeta_bar_H
  TB$L_zeta_bar[iter] = equi_sol$zeta_bar_L
  TB$r1[iter] = equi_sol$result_qH$interest_rate_fun(zeta_cases[1])
  TB$r2[iter] = equi_sol$result_qH$interest_rate_fun(zeta_cases[2])
  TB$d_hat1[iter] = equi_sol$result_qL$endogenous_private_debt_limit_fun(zeta_cases[1])
  TB$d_hat2[iter] = equi_sol$result_qL$endogenous_private_debt_limit_fun(zeta_cases[2])
  TB$xL1[iter] = equi_sol$result_qL$l_fun(zeta_cases[1])
  TB$xL2[iter] = equi_sol$result_qL$l_fun(zeta_cases[2])
  TB$xH1[iter] = equi_sol$result_qH$l_fun(zeta_cases[1])
  TB$xH2[iter] = equi_sol$result_qH$l_fun(zeta_cases[2])
  TB$output_change[iter] = Delta_output_fun(equi_sol$kappaH, equi_sol$kappaL, w0)
}
pdf(paste0("../Figures/FigureA7a_govt_financing_and_DeltaK.pdf"), width=plot_width, height=plot_height, family="serif")
par(mar = margins, mfrow=c(1,1))
MyPlot(TB$g_bar, TB[,list(dK*100)], xlab=TeX("govt program size $\\bar{g}$"), ylab=TeX("capital quantity change $\\Delta K$  (%)")); abline(h=0)
dev.off()
pdf(paste0("../Figures/FigureA7b_govt_financing_and_Deltaw.pdf"), width=plot_width, height=plot_height, family="serif")
par(mar = margins, mfrow=c(1,1))
MyPlot(TB$g_bar, TB$dw*100, xlab=TeX("govt program size $\\bar{g}$"), ylab=TeX("capital quality change  $\\Delta \\omega$ (%)") ); abline(h=0)
dev.off()


# ----- Figure A8: Credit Intervention and Welfare ----------
# Solve for the Welfare function
g_bar_vec = seq(0,AH,length.out=30)
w_grid = seq(0.0001,0.9999,length.out = 500)
init_vec = rep(0, length(g_bar_vec))
index = length(w_grid)/2+1
TB = data.table( g_bar=g_bar_vec, g_takeup_avg=init_vec, g_takeup=init_vec, g_takeup2=init_vec, g_takeup3=init_vec, productivity_avg=init_vec, welfare_avg=init_vec, welfare = init_vec, welfare2 = init_vec, welfare3=init_vec, l_dN=init_vec, w0=init_vec, muw=init_vec, muK=init_vec, Delta_w=init_vec, Delta_K=init_vec, qH=init_vec, qL=init_vec )
for(iter in 1:length(g_bar_vec)){
  print( paste0("Progress: ", round(100*iter/length(g_bar_vec),1), "%") )
  result_welfare = solve_welfare(w_grid, d_bar, g_bar=g_bar_vec[iter])
  # Another idea: compare the welfare with and without the slippery slope of credit intervention. 
  equi_sol = solve_other_moments(result_welfare$equi_sol)
  w0 = equi_sol$w_avg
  TB$welfare[iter] = result_welfare$W_fun(0.01)
  TB$welfare2[iter] = result_welfare$W_fun(0.5)
  TB$welfare3[iter] = result_welfare$W_fun(1)
  TB$welfare_avg[iter] = result_welfare$W_fun(w0)
  TB$g_takeup[iter] = equi_sol$total_take_up_fun(0.01)
  TB$g_takeup2[iter] = equi_sol$total_take_up_fun(0.5)
  TB$g_takeup3[iter] = equi_sol$total_take_up_fun(0.99)
  TB$g_takeup_avg[iter] = equi_sol$total_takeup
  TB$productivity_avg[iter] = equi_sol$avg_productivity
  TB$l_dN[iter] = interp1( result_welfare$w_grid, result_welfare$l_dN, w0 )
  TB$w0[iter] = w0
  TB$muw[iter]  = interp1( result_welfare$w_grid, result_welfare$muw, w0 )
  TB$muK[iter] = interp1( result_welfare$w_grid, result_welfare$muK, w0 ) 
  TB$Delta_w[iter] = interp1( result_welfare$w_grid, result_welfare$Delta_w, w0 )
  TB$Delta_K[iter] = interp1( result_welfare$w_grid, result_welfare$Delta_K, w0 ) 
  TB$qH[iter] = result_welfare$qH
  TB$qL[iter] = result_welfare$qL
}
index = which(TB$welfare_avg==max(TB$welfare_avg))
pdf(paste0("../Figures/FigureA8a_w_avg_and_g_bar.pdf"), width=plot_width, height=plot_height, family="serif")
par(mfrow=c(1,1), mar=margins)
MyPlot(TB$g_bar, TB[,list(w0)], xlab=TeX("govt program size $\\bar{g}$"), ylab=TeX("average capital quality") )
dev.off()
pdf(paste0("../Figures/FigureA8b_welfare_and_g_bar.pdf"), width=plot_width, height=plot_height, family="serif")
par(mfrow=c(1,1), mar=margins)
MyPlot(TB$g_bar, TB[,list(100*(welfare_avg/welfare_avg[1]-1))], xlab=TeX("govt program size $\\bar{g}$"), ylab=TeX("welfare (% diff from laissez-faire)")  )
abline(v=TB$g_bar[index], lty=3)
dev.off()

# -------------------- Figure A9: Slippery Slope in the Same Economy --------------------
# Government intervention and the avg w
g_bar_vec = c( seq(0, g_bar ,length.out=20), seq(g_bar+0.005, 0.25, length.out=20) )
index_baseline = which.min( abs(g_bar_vec-0.14)  )
# Solve for the GDP_drop target to achieve the same as the baseline when g_bar=0
resultH = post_crisis_fraction_of_capital(equi_sol_no_govt$qH, d_bar, 0, type="H")
resultL = post_crisis_fraction_of_capital(equi_sol_no_govt$qL, d_bar, 0, type="L")
kappaH = resultH$kappa
kappaL = resultL$kappa
w = w0  + Deltaw_fun(kappaH, kappaL, w0) # immediately after the crisis
K = K0 * (1 + DeltaK_fun(kappaH, kappaL, w0) )
for(iter in 1:10/dt){
  # 10 years
  w = w + muw_fun(equi_sol_no_govt$iH, equi_sol_no_govt$iL, w)*dt
  K = K * (1 + muK_fun(equi_sol_no_govt$iH, equi_sol_no_govt$iL, w)*dt )
}
GDP_drop_target = output_drop_and_gbar(0.096, w, equi_sol_no_govt$qH, equi_sol_no_govt$qL)

# intervention and next intervention scale. 
govt_funding_next_crisis_vec = rep(0, length(g_bar_vec))
equi_sol_no_govt = solve_equilibrium(d_bar, g_bar=0)
equi_sol_no_govt = solve_other_moments(equi_sol_no_govt)
w0 = equi_sol_no_govt$w_avg
K0 = 1
# Think about the actual amount of intervention. 
govt_credit_provision_vec = g_bar_vec
govt_credit_provision_next_crisis = g_bar_vec
w_next_crisis_vec = g_bar_vec
K_next_crisis_vec = g_bar_vec
productivity_vec = rep(0,length(g_bar_vec))
productivity_next_crisis_vec = rep(0,length(g_bar_vec))
for(iter_g_bar in 1:length(g_bar_vec)){
  print( paste0("Progress:", iter_g_bar/length(g_bar_vec)*100, "%") )
  resultH = post_crisis_fraction_of_capital(equi_sol_no_govt$qH, d_bar, g_bar_vec[iter_g_bar], type="H")
  resultL = post_crisis_fraction_of_capital(equi_sol_no_govt$qL, d_bar, g_bar_vec[iter_g_bar], type="L")
  kappaH = resultH$kappa
  kappaL = resultL$kappa
  govt_credit_provision_vec[iter_g_bar] = w0*resultH$govt_sector_financing + (1-w0)*resultL$govt_sector_financing
  productivity_vec[iter_g_bar] = w0*AH + (1-w0)*AL
  w = w0  + Deltaw_fun(kappaH, kappaL, w0) # immediately after the crisis
  K = K0 * (1 + DeltaK_fun(kappaH, kappaL, w0) )
  for(iter in 1:10/dt){
    # 10 years
    w = w + muw_fun(equi_sol_no_govt$iH, equi_sol_no_govt$iL, w)*dt
    K = K * (1 + muK_fun(equi_sol_no_govt$iH, equi_sol_no_govt$iL, w)*dt )
  }  
  w_next_crisis_vec[iter_g_bar] = w
  K_next_crisis_vec[iter_g_bar] = K
  result = uniroot( function(g_bar) output_drop_and_gbar(g_bar, w, equi_sol_no_govt$qH, equi_sol_no_govt$qL) - GDP_drop_target, interval = c(0, 2)  )
  govt_funding_next_crisis_vec[iter_g_bar] = result$root
  
  resultH= post_crisis_fraction_of_capital(equi_sol_no_govt$qH, d_bar, govt_funding_next_crisis_vec[iter_g_bar], type="H")
  resultL= post_crisis_fraction_of_capital(equi_sol_no_govt$qL, d_bar, govt_funding_next_crisis_vec[iter_g_bar], type="L")
  # Total credit provision
  govt_credit_provision_next_crisis[iter_g_bar] = w*resultH$govt_sector_financing + (1-w)*resultL$govt_sector_financing
  productivity_next_crisis_vec[iter_g_bar] = w*AH + (1-w)*AL
}
sensitivity = ( govt_funding_next_crisis_vec[index_baseline] - govt_funding_next_crisis_vec[1] ) / ( g_bar_vec[index_baseline] - g_bar_vec[1] )
pdf(paste0("../Figures/FigureA9a_same_economy_slippery_slope_w.pdf"), width=plot_width,height=plot_height, family="serif")
par(mfrow=c(1,1), mar=margins)
MyPlot( g_bar_vec, w_next_crisis_vec, xlab=TeX("current govt program size $\\bar{g}$"), ylab=TeX("next crisis (in 10 years) $\\omega$'") )
dev.off()
pdf(paste0("../Figures/FigureA9b_same_economy_slippery_slope_credit.pdf"), width=plot_width,height=plot_height, family = "serif")
par(mfrow=c(1,1), mar=margins)
MyPlot( g_bar_vec, govt_funding_next_crisis_vec, xlab=TeX("current govt program size $\\bar{g}$"), ylab=TeX("next crisis (in 10 years) govt intervention $\\bar{g}'$") )
dev.off()
